#ifndef YOLO_INTERFACE_H__
#define YOLO_INTERFACE_H__
#include "ingenic_mnni.h"
#include "common/common_utils.h"
#include "utils/imgproc.h"
#include "utils/postproc.h"
#include "common/common_utils.h"
#include "data_type.h"

using namespace magik::uranus;
using namespace std;

namespace Detect{
    /**
     * @brief YOLOV5检测配置参数结构体
     * 包含置信度阈值、NMS阈值、模型路径等YOLOV5相关配置
     */
    typedef struct {
        float conf_thresh = 0.5;    ///< 目标检测置信度阈值，高于此值的检测框被保留
        float nms_thres = 0.5;      ///< NMS非极大值抑制阈值
        int classes = 1;            ///< 检测类别数量
        int box_num = 3;            ///< 每个网格单元的预测框数量
        std::string model_path;     ///< 模型文件路径
        int input_w = 640;          ///< 模型输入图像宽度
        int input_h = 640;          ///< 模型输入图像高度
        magik::uranus::ChannelLayout image_fmt = magik::uranus::ChannelLayout::RGBA;  ///< 输入图像通道格式
    }YOLOV5Config;

    /**
     * @brief YOLOV5目标检测接口类
     * 提供基于YOLOV5算法的目标检测功能，包括模型加载、推理和后处理等
     */
    class YOLOV5Interface {
        private:
            std::vector<float> strides = {8.0, 16.0, 32.0};  ///< 特征图下采样步长
            std::vector<float> anchor = {10, 13, 16, 30, 33, 23,  ///< 锚框尺寸
                                    30, 61, 62, 45, 59, 119, 
                                        116, 90, 156, 198, 373, 326};
            DetectorType_t detector_type = DetectorType_t::YOLOV5;  ///< 检测器类型
            Detect::YOLOV5Config  yolo5_config;  ///< YOLOV5配置参数
            magik::uranus::Network* model=nullptr;  ///< 推理网络指针

            /**
             * @brief 生成候选边界框
             * @param out_res 模型输出结果
             * @param candidate_boxes 候选边界框列表
             * @param img_w 输入图像宽度
             * @param img_h 输入图像高度
             */
            void genBBox(vector<Tensor> out_res, 
                std::vector<ObjBbox_t>& candidate_boxes, 
                int img_w, 
                int img_h
            );

            /**
             * @brief 缩放边界框到原始图像尺寸
             * @param ori_w 原始图像宽度
             * @param ori_h 原始图像高度
             * @param new_w 缩放后图像宽度
             * @param new_h 缩放后图像高度
             * @param out_boxes 边界框列表
             */
            void scale_boxes(int ori_w, int ori_h, int new_w, int new_h, 
                std::vector<magik::uranus::ObjBbox_t> &out_boxes);

            /**
             * @brief 运行模型推理
             * @param model 网络模型指针
             * @param img 输入图像Tensor
             * @param outputs 模型输出结果
             */
            void run_model(Network *model, Tensor &img, vector<Tensor> &outputs);

        public:
            /**
             * @brief 构造函数
             * @param config YOLOV5配置参数
             */
            YOLOV5Interface(Detect::YOLOV5Config config);

            /**
             * @brief 加载YOLOV5模型
             * @return 0表示成功，非0表示失败
             */
            int loadModel();

            /**
             * @brief 释放模型资源
             * @return 0表示成功，非0表示失败
             */
            int releaseModel();

            /**
             * @brief 打印模型前向传播内存使用情况
             */
            void print_model_forward_memory();

            /**
             * @brief 执行目标检测
             * @param working_img 输入图像Tensor
             * @param orgin_width 原始图像宽度
             * @param orgin_height 原始图像高度
             * @param output_boxes 检测结果边界框列表
             */
            void detect(magik::uranus::Tensor working_img, int orgin_width, int orgin_height, vector<ObjBbox_t>& output_boxes);
        };

} // namespace Detect
#endif